from typing import Tuple, List
from nlp_tools import macros as k
import os
CORPUS_PATH = os.path.join(k.DATA_PATH, 'corpus')

import pandas as pd
from sklearn.model_selection import train_test_split
from nlp_tools.utils.data import unison_shuffled_copies


class BaseCorpusLoader():
    __name__ = "base_corpus_loader"

    def __init__(self, data_path, shuffle: bool = True, split_train_test=False, split_test_rate=0.1, max_length=128):
        '''
            加载训练数据
            :param data_path: 数据路径
            :param shuffle: 是否打乱顺序
            :param split_train_test: 是否分割训练集和测试集
            :param split_test_rate: 测试集占比
            :return:
            '''
        self.data_path = data_path
        self.shuffle = shuffle
        self.split_train_test = split_train_test
        self.split_test_rate = split_test_rate
        self.max_length = max_length

    def load_data(self, shuffle: bool = None):
        if shuffle == None:
            shuffle = self.shuffle

        if type(self.data_path) == str:
            df = pd.read_csv(self.data_path)
        else:
            df = self.data_path

        x, y = self.get_x_and_y(df)

        if shuffle:
            x, y = unison_shuffled_copies(x, y)

        data_union = [(x, y) for x, y in zip(x, y)]
        if self.split_train_test:
            return train_test_split(data_union, test_size=self.split_test_rate)
        else:
            return data_union

    def get_x_and_y(self, df):
        '''
        处理dataframe对象，返回 对应的输入x和标签 y
        :param df:
        :return:
        '''
        raise NotImplementedError("not implemment yeat~")



class DataReader(object):

    @staticmethod
    def read_conll_format_file(file_path: str,
                               text_index: int = 0,
                               label_index: int = 1) -> Tuple[List[List[str]], List[List[str]]]:
        """
        Read conll format data_file
        Args:
            file_path: path of target file
            text_index: index of text data, default 0
            label_index: index of label data, default 1

        Returns:

        """
        x_data, y_data = [], []
        with open(file_path, 'r', encoding='utf-8') as f:
            lines = f.read().splitlines()
            x, y = [], []
            for line in lines:
                rows = line.split(' ')
                if len(rows) == 1:
                    x_data.append(x)
                    y_data.append(y)
                    x = []
                    y = []
                else:
                    x.append(rows[text_index])
                    y.append(rows[label_index])
        return x_data, y_data

from nlp_tools.corpus.ner.corpus_loader import ChineseDailyNerCorpus
from nlp_tools.corpus.classify.corpus_loader import SMP2018ECDTCorpus
from nlp_tools.corpus.classify.dialogue_corpus_loader import DialogueCorpusLoader